import json
import re

from elasticsearch import Elasticsearch

from message_management.case_utils.document_law import search_law, LawDocument
from utils.model_api import ModelAPI

es = Elasticsearch(["10.206.60.14:9200"])

# 预定义的法典名称列表，包括一些常见变体
law_names = ["刑法", "民法典", "商法", "行政法", "宪法"]
llm = ModelAPI()


def construct_class(text):
    # 去除前后空格、换行符等
    text = text.strip()


    # todo 把检索出的数据再进行处理、检索。如果**条存在什么的


    pattern_map = {
        "volume": re.compile(r"(第[\d一二三四五六七八九十百零]+编)"),
        "subvolume": re.compile(r"(第[\d一二三四五六七八九十百零]+分编)"),
        "chapter": re.compile(r"(第[\d一二三四五六七八九十百零]+章)"),
        "section": re.compile(r"(第[\d一二三四五六七八九十百零]+节)"),
        "article_number": re.compile(r"(第[\d一二三四五六七八九十百千零]+条)")
    }
    context = {"law_name": None, "volume": None, "subvolume": None, "chapter": None, "section": None,
               "article_number": None, "article_text": None}
    for key, pattern in pattern_map.items():
        match = pattern.search(text)
        if match:
            context[key] = match.group(1).strip()
    found_laws = find_law_names(text, law_names)
    # print(found_laws)
    # print("====")
    if found_laws and len(found_laws) != 0:
        context["law_name"] = " ".join(found_laws)
    context["article_text"] = text

    return context

def construct_es_query(params):
    must_clauses = []
    should_clauses = []

    if params.get("law_name") is not None:
        must_clauses.append(f'{{"match": {{"law_name": "{params["law_name"]}"}}}}')

    if params.get("volume") is not None:
        must_clauses.append(
            f'{{"match_phrase_prefix": {{"volume": {{"query": "{params["volume"]}", "max_expansions": 10}}}}}}')

    if params.get("subvolume") is not None:
        must_clauses.append(
            f'{{"match_phrase_prefix": {{"subvolume": {{"query": "{params["subvolume"]}", "max_expansions": 10}}}}}}')

    if params.get("chapter") is not None:
        must_clauses.append(
            f'{{"match_phrase_prefix": {{"chapter": {{"query": "{params["chapter"]}", "max_expansions": 10}}}}}}')

    if params.get("section") is not None:
        must_clauses.append(
            f'{{"match_phrase_prefix": {{"section": {{"query": "{params["section"]}", "max_expansions": 10}}}}}}')

    if params.get("article_number") is not None:
        must_clauses.append(f'{{"term": {{"article_number": "{params["article_number"]}"}}}}')

    # article_text is added to the 'should' clause
    if params.get("article_text") is not None:
        should_clauses.append(f'{{"match": {{"article_text": "{params["article_text"]}"}}}}')

    if should_clauses:
        query = (
            f'{{"query": {{"bool": {{"must": [{", ".join(must_clauses)}], "should": [{", ".join(should_clauses)}]}}}}}}'
        )
    else:
        query = f'{{"query": {{"bool": {{"must": [{", ".join(must_clauses)}]}}}}}}'

    return query


def create_prompt(law_data, user_question):
    # 初始化法律信息列表
    law_info_list = []

    # 构建法律信息对象
    for document in law_data:
        law_info = {
            "law_name": document.law_name,
            "volume": document.volume if document.volume else "",
            "chapter": document.chapter if document.chapter else "",
            "article_number": document.article_number if document.article_number else "",
            "article_text": document.article_text if document.article_text else ""
        }
        # 添加到列表
        law_info_list.append(law_info)

    # 构建完整的JSON对象
    prompt_data = {
        "法条数据": law_info_list,
        "question": user_question,
        "注意": "根据question字段中的问题，结合法条数据进行问答。不要回答question字段中问题之外的法条数据。"
    }

    # 将对象转换为JSON字符串
    prompt_json = json.dumps(prompt_data, ensure_ascii=False, indent=4)
    return prompt_json


def find_law_names(text, law_names):
    # 修改正则表达式，移除 \b 边界，以适应更多场景
    pattern = re.compile(r'(' + '|'.join(map(re.escape, law_names)) + r')', re.IGNORECASE)

    # 搜索文本并找到所有匹配的法典名称
    matches = pattern.findall(text)

    # 打印用于调试的正则表达式和匹配结果
    # print("Regex Pattern:", pattern.pattern)
    # print("Matches Found:", matches)

    # 返回去重后的匹配列表
    return list(set(matches))


def law_chat(text):
    params = construct_class(text)
    # print(params)
    es_query = construct_es_query(params)
    print(es_query)
    data, total = search_law(es_query)
    # 如果超过5个结果，只显示前5个
    if total > 8:
        data = data[:8]
    prompt = create_prompt(data, text)
    print("=================")
    print(prompt)
    # llm_response = llm.chat(prompt)
    for data in llm.chat_stream(prompt):
        yield data  # 直接传递 run_glm 生成的每个字节流项


if __name__ == '__main__':
    res = law_chat("【过失犯罪】应当预见自己的行为可能发生危害社会的结果，是第几章第几节第几条？")
    for r in res:
        print(r)
    # example_text = "【过失犯罪】应当预见自己的行为可能发生危害社会的结果，是第几章第几节第几条？"
    # # example_text = "本案涉及民法典和刑法的应用。商法也可能是相关的，但没有涉及行政法或宪法。"
    #
    # params = construct_class(example_text)
    # # print(params)
    # es_query = construct_es_query(params)
    # print(es_query)
    # data, total = search_law(es_query)
    # # 如果超过5个结果，只显示前5个
    # if total > 5:
    #     data = data[:5]
    #
    # prompt = create_prompt(data, example_text)
    # print("=================")
    # print(prompt)
    # llm_response = llm.chat(prompt)
